package com.cheil.publish.service.ai;

import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesis;
import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesisParam;
import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesisResult;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.fastjson2.JSON;
import java.util.List;
import lombok.extern.log4j.Log4j2;
import lombok.val;
import org.springframework.ai.image.Image;
import org.springframework.ai.image.ImageGeneration;
import org.springframework.ai.image.ImageGenerationMetadata;
import org.springframework.ai.image.ImageModel;
import org.springframework.ai.image.ImageOptions;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;

/**
 * DashscopeImageModel是实现ImageModel接口的类，用于调用阿里云Dashscope文生图API。
 */
@Log4j2
public class DashscopeImageModel implements ImageModel {

    private final String apiKey;
    private final DashscopeImageOptions defaultOptions;
    private final RetryTemplate retryTemplate;
    private final ImageSynthesis imageSynthesis;

    public DashscopeImageModel(String apiKey) {
        this(apiKey, DashscopeImageOptions.builder().build(), RetryUtils.DEFAULT_RETRY_TEMPLATE);
    }

    public DashscopeImageModel(String apiKey, DashscopeImageOptions defaultOptions) {
        this(apiKey, defaultOptions, RetryUtils.DEFAULT_RETRY_TEMPLATE);
    }

    public DashscopeImageModel(String apiKey, DashscopeImageOptions defaultOptions, RetryTemplate retryTemplate) {
        Assert.hasText(apiKey, "Dashscope API key must not be empty");
        Assert.notNull(defaultOptions, "defaultOptions must not be null");
        Assert.notNull(retryTemplate, "retryTemplate must not be null");

        this.apiKey = apiKey;
        this.defaultOptions = defaultOptions;
        this.retryTemplate = retryTemplate;
        this.imageSynthesis = new ImageSynthesis();
    }

    @Override
    public ImageResponse call(ImagePrompt imagePrompt) {
        return this.retryTemplate.execute(ctx -> {
            try {
                ImageSynthesisParam imageRequest = createRequest(imagePrompt);
                ImageSynthesisResult imageResult = this.imageSynthesis.call(imageRequest);
                return convertResponse(imageResult);
            } catch (ApiException | NoApiKeyException e) {
                throw new RuntimeException(e);
            }
        });
    }

    private ImageSynthesisParam createRequest(ImagePrompt imagePrompt) {
        String prompt = imagePrompt.getInstructions().get(0).getText();

        // Merge runtime options with default options
        DashscopeImageOptions runtimeOptions = null;
        if (imagePrompt.getOptions() != null) {
            runtimeOptions = ModelOptionsUtils.copyToTarget(imagePrompt.getOptions(), ImageOptions.class,
                    DashscopeImageOptions.class);
        }

        DashscopeImageOptions options = mergeOptions(runtimeOptions, this.defaultOptions);

        val paramBuilder = ImageSynthesisParam.builder()
                                              .apiKey(this.apiKey)
                                              .model(options.getModel())
                                              .prompt(prompt)
                                              .n(options.getN())
                                              .size(options.getSize());

        if (options.getNegativePrompt() != null) {
            paramBuilder.negativePrompt(options.getNegativePrompt());
        }


        if (options.getSeed() != null) {
            try {
                paramBuilder.seed(Integer.valueOf(options.getSeed()));
            } catch (NumberFormatException e) {
                log.warn("Invalid seed value: {}", options.getSeed());
            }
        }

        return paramBuilder.build();
    }

    private DashscopeImageOptions mergeOptions(DashscopeImageOptions runtimeOptions, DashscopeImageOptions defaultOptions) {
        if (runtimeOptions == null) {
            return defaultOptions;
        }

        DashscopeImageOptions.Builder builder = DashscopeImageOptions.builder()
                .model(ModelOptionsUtils.mergeOption(runtimeOptions.getModel(), defaultOptions.getModel()))
                .n(ModelOptionsUtils.mergeOption(runtimeOptions.getN(), defaultOptions.getN()))
                .size(ModelOptionsUtils.mergeOption(runtimeOptions.getSize(), defaultOptions.getSize()))
                .negativePrompt(ModelOptionsUtils.mergeOption(runtimeOptions.getNegativePrompt(), defaultOptions.getNegativePrompt()))
                .promptExtend(ModelOptionsUtils.mergeOption(runtimeOptions.getPromptExtend(), defaultOptions.getPromptExtend()))
                .watermark(ModelOptionsUtils.mergeOption(runtimeOptions.getWatermark(), defaultOptions.getWatermark()))
                .seed(ModelOptionsUtils.mergeOption(runtimeOptions.getSeed(), defaultOptions.getSeed()));

        return builder.build();
    }

    private ImageResponse convertResponse(ImageSynthesisResult imageResult) {
        log.info("{}" , JSON.toJSONString(imageResult));
        if (imageResult == null || imageResult.getOutput() == null || imageResult.getOutput().getResults() == null) {
            return new ImageResponse(List.of());
        }

        List<ImageGeneration> imageGenerationList = imageResult.getOutput().getResults().stream()
                .map(result -> {
                    String imageUrl = result.getOrDefault("url","");
                    String orig_prompt = result.getOrDefault("orig_prompt","");
                    String actual_prompt = result.getOrDefault("actual_prompt","");
                    Image image = new Image(imageUrl,null);
                    ImageGenerationMetadata metadata = new DashscopeImageGenerationMetadata(actual_prompt);
                    return new ImageGeneration(image, metadata);
                })
                .toList();

        return new ImageResponse(imageGenerationList);
    }
}